#include "MixTrack.hpp"
#include <vector>
#include <Eigen/Core>
#include <Eigen/Dense>
#include <opencv2/core/eigen.hpp>
#include "trt_common/trt_infer.hpp"
#include "trt_common/ilogger.hpp"

namespace MixTrack {
    using namespace std;
    struct Config {
        // std::vector<float> window;
        float template_factor = 2.0;
        float search_factor = 4.5; // 5.0
        int template_size = 112; //192
        int search_size = 224; // 384
        int stride = 16;
        int feat_sz = 14; // 24
        int update_interval = 200;
    };

    class TrackerImpl : public Tracker{
    public:
        ~TrackerImpl() = default;

        bool startup(const std::string &engine_path, int gpuid){
            gpu_ = gpuid;
            TRT::set_device(gpuid);

            infer_model_ = TRT::load_infer(engine_path);
            if(infer_model_ == nullptr){
                INFOE("Load model failed: %s", engine_path.c_str());
                return false;
            }
            stream_ = infer_model_->get_stream();
            infer_model_->print();

            return true;
        }

        void init(cv::Mat &img, cv::Rect &init_bbox) override{
            frame_width = img.cols;
            frame_hegith = img.rows;

            // get subwindow
            cv::Mat zt_patch; 
            cv::Mat ozt_patch;   
            float resize_factor = 1.f;
            this->sample_target(img, zt_patch, init_bbox, this->cfg.template_factor, this->cfg.template_size, resize_factor);
            this->sample_target(img, ozt_patch, init_bbox, this->cfg.template_factor, this->cfg.template_size, resize_factor);
            // cv::Mat oz_patch = z_patch.clone(); 
            this->z_patch = zt_patch;
            this->oz_patch = ozt_patch;
            this->state = init_bbox;

            std::cout << "init bbox:" << state << std::endl;
        }

        cv::Rect track(cv::Mat& img) override{
            cv::Mat x_patch;
            this->frame_id += 1;
            float resize_factor = 1.f;
            int ret = this->sample_target(img, x_patch, this->state, this->cfg.search_factor, this->cfg.search_size, resize_factor);
            if(ret == -1)
            {
                printf("invalid box! \n");
                return cv::Rect(-1, -1, -1, -1);
            }
            
            // preprocess input tensor
            float m[] = {0.406, 0.456, 0.485};
            float std[]  = {0.225, 0.224, 0.229};

            zin_ = infer_model_->input(0);
            zin_->set_norm_mat_invert(0, z_patch, m, std);

            ozin_ = infer_model_->input(1);
            ozin_->set_norm_mat_invert(0, oz_patch, m, std);
            
            xin_ = infer_model_->input(2);
            xin_->set_norm_mat_invert(0, x_patch, m, std);

            boxes_map_ = infer_model_->output(0);
            score_map_ = infer_model_->output(1);

            // 推理
            infer_model_->forward();

            float *output_pred_boxes = boxes_map_->cpu<float>();
            float *output_pred_scores = score_map_->cpu<float>();

            cv::Rect2f pred_box;
            float pred_score;

            // printf("output_pred_boxes:%f %f %f %f, score:%f \n", output_pred_boxes[0], output_pred_boxes[1], output_pred_boxes[2], output_pred_boxes[3], output_pred_scores[0]);
            
            this->cal_bbox(output_pred_boxes, output_pred_scores, pred_box, pred_score, resize_factor);
            this->map_box_back(pred_box, resize_factor);
            this->clip_box(pred_box, img.rows, img.cols, 10);
            
            this->state = pred_box;

            this->max_pred_score = this->max_pred_score * this->max_score_decay;
            // update template
            if (pred_score > 0.9 && pred_score > this->max_pred_score)
            {
                ret = this->sample_target(img, this->max_oz_patch, this->state, this->cfg.template_factor, this->cfg.template_size, resize_factor);
                if(ret == -1)
                {
                    return cv::Rect(-1, -1, -1, -1);
                }
                
                this->max_pred_score = pred_score;
            }

            if (this->frame_id % this->cfg.update_interval == 0)
            {
                this->oz_patch = this->max_oz_patch;
                this->max_pred_score = -1.0;
                this->max_oz_patch = this->oz_patch;
            }

            if (pred_score < 0.7 && this->frame_id > 30)
            {
                printf("pred score:%f \n", pred_score);
                return cv::Rect(-1, -1, -1, -1);
            }
            
            return pred_box;
        }
    private:

        void cal_bbox(float *boxes_ptr, float * scores_ptr, cv::Rect2f &pred_box, float &max_score, float resize_factor) {
            auto cx = boxes_ptr[0];
            auto cy = boxes_ptr[1];
            auto w = boxes_ptr[2];
            auto h = boxes_ptr[3];
            // std::cout << "cal_bbox cx cy w h "<< cx << " " << cy << " " << w << " " << h << std::endl;
            cx = cx * this->cfg.search_size / resize_factor;
            cy = cy * this->cfg.search_size / resize_factor;
            w = w * this->cfg.search_size / resize_factor;
            h = h * this->cfg.search_size / resize_factor;
            
            pred_box.x = cx - 0.5 * w;
            pred_box.y = cy - 0.5 * h;
            pred_box.width = w;
            pred_box.height = h;

            max_score = scores_ptr[0];

        }

        void map_box_back(cv::Rect2f &pred_box, float resize_factor) {
            float cx_prev = this->state.x + 0.5 * (this->state.width);
            float cy_prev = this->state.y + 0.5 * (this->state.height);

            float half_side = 0.5 * this->cfg.search_size / resize_factor;

            float w = pred_box.width;
            float h = pred_box.height;
            float cx = pred_box.x + 0.5 * w;
            float cy = pred_box.y + 0.5 * h;

            float cx_real = cx + (cx_prev - half_side);
            float cy_real = cy + (cy_prev - half_side);

            pred_box.x = cx_real - 0.5 * w;
            pred_box.y = cy_real - 0.5 * h;
            pred_box.width = w;
            pred_box.height = h;
        }
        void clip_box(cv::Rect2f &box, int height, int wight, int margin) {
            box.x = std::min(std::max(0, int(box.x)), wight - margin);
            box.y = std::min(std::max(0, int(box.y)), height - margin);
            box.width = std::min(std::max(margin, int(box.x + box.width)), wight) - box.x;
            box.height= std::min(std::max(margin, int(box.y + box.height)), height) - box.y;
        }

        int  sample_target(const cv::Mat &im, cv::Mat &croped, cv::Rect2f target_bb, float search_area_factor, int output_sz, float &resize_factor) {
            /* Extracts a square crop centrered at target_bb box, of are search_area_factor^2 times target_bb area

            args:
                im: Img image
                target_bb - target box [x0, y0, x1, y1]
                search_area_factor - Ratio of crop size to target size
                output_sz - Size
            
            */
            cv::Rect2f bbox = target_bb;
            bbox &= cv::Rect2f(0, 0, frame_width, frame_hegith);

            if(bbox.width <= 0 || bbox.height <= 0 )
            {
                std::cout << "invalid bbox:" << bbox << std::endl;
                return -1;
            }
            
            int x = bbox.x;
            int y = bbox.y;
            int w = bbox.width;
            int h = bbox.height;

            int crop_sz = std::ceil(std::sqrt(w *h) * search_area_factor);

            float cx = x + 0.5 * w;
            float cy = y + 0.5 * h;
            int x1 = std::round(cx - crop_sz * 0.5);
            int y1 = std::round(cy - crop_sz * 0.5);

            int x2 = x1 + crop_sz;
            int y2 = y1 + crop_sz;

            int x1_pad = std::max(0, -x1);
            int x2_pad = std::max(x2 - im.cols +1, 0);
            
            int y1_pad = std::max(0, -y1);
            int y2_pad = std::max(y2- im.rows + 1, 0);

            // Crop target
            cv::Rect roi_rect(x1+x1_pad, y1+y1_pad, (x2-x2_pad)-(x1+x1_pad), (y2-y2_pad)-(y1+y1_pad));
            cv::Mat roi = im(roi_rect);

            // Pad
            cv::copyMakeBorder(roi, croped, y1_pad, y2_pad, x1_pad, x2_pad, cv::BORDER_CONSTANT);

            // Resize
            cv::resize(croped, croped, cv::Size(output_sz, output_sz));

            resize_factor = output_sz * 1.f / crop_sz;

            return 0;
        }

        int gpu_ = 0;
        shared_ptr<TRT::Infer> infer_model_;
        TRT::CUStream stream_ = nullptr;

        int frame_width = 0;
        int frame_hegith = 0;

        cv::Mat z_patch; // template
        cv::Mat oz_patch; // online_template
        cv::Mat max_oz_patch; // online max template

        // config static
        Config cfg; 
        cv::Rect2f state;

        // init dynamic input dims
        std::vector<std::vector<int64_t>> input_node_dims = {
            {1, 3, 112, 112}, // z  (b=1,c,h,w)
            {1, 3, 112, 112}, // z  (b=1,c,h,w)
            {1, 3, 224, 224} // x
        };

        // Define FP32 mean and scale values
        const float mean_vals[3] = {0.485f * 255.f, 0.456f * 255.f, 0.406f * 255.f};  // RGB
        const float norm_vals[3] = {1 / 0.229f / 255.f, 1 / 0.224f / 255.f, 1 / 0.225f / 255.f};

        float max_pred_score = -1.f;
        float max_score_decay = 1.f;
        int update_count = 0;

        int frame_id = 0;
        shared_ptr<TRT::Tensor>  zin_;
        shared_ptr<TRT::Tensor>  ozin_;
        shared_ptr<TRT::Tensor>  xin_;
        
        shared_ptr<TRT::Tensor> score_map_;
        shared_ptr<TRT::Tensor> boxes_map_;

        
    };

    shared_ptr<Tracker> create_tracker(const std::string &engine_path,int gpuid){
        shared_ptr<TrackerImpl> instance(new TrackerImpl{});
        if(!instance->startup(engine_path, gpuid))
            instance.reset();
        return instance;
    }
}